// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <iostream>

#include <pollux/vector/flat_vector.h>
#include <pollux/vector/fuzzer/generator_spec.h>
#include <tests/vector/fuzzer/examples/Utils.h>

using namespace kumo::pollux;
using namespace generator_spec_maker;
using namespace generator_spec_examples;

int main() {
    // This example shows how to use our GeneratorSpec class to generate vectors
    // of scalars with user defined distributions.

    constexpr int sampleSize = 100000;
    constexpr int32_t lo = 0, hi = 10;
    constexpr double mu = 5.0, sigma = 2.0;
    constexpr double userLo = 0.01, userHi = 0.99;
    constexpr double nullProbability = 0.18;

    auto uniform = std::uniform_int_distribution<int32_t>(lo, hi);
    auto normal = std::normal_distribution<double>(mu, sigma);
    auto coinToss = std::uniform_real_distribution<double>();
    auto userDefined =
            [userLo, userHi, &coinToss](FuzzerGenerator &rng) -> double {
        auto val = coinToss(rng);
        if (val <= userLo) {
            return val;
        } else if (val > userHi) {
            return val * val;
        } else {
            return 0.0;
        }
    };

    FuzzerGenerator rng;
    memory::MemoryManager::initialize({});
    auto pool = memory::memoryManager()->addLeafPool();

    GeneratorSpecPtr uniformIntGenerator =
            RANDOM_INTEGER(uniform, nullProbability);
    GeneratorSpecPtr normalDoubleGenerator =
            RANDOM_DOUBLE(normal, nullProbability);
    GeneratorSpecPtr userDefinedGenerator =
            RANDOM_DOUBLE(userDefined, nullProbability);

    VectorPtr uniformVector =
            uniformIntGenerator->generateData(rng, pool.get(), sampleSize);
    VectorPtr normalVector =
            normalDoubleGenerator->generateData(rng, pool.get(), sampleSize);
    VectorPtr userDefinedVector =
            userDefinedGenerator->generateData(rng, pool.get(), sampleSize);

    std::cout << "Sample data generated from uniform distribution:\n"
            << plotVector(uniformVector->as_flat_vector<int32_t>()) << "\n";

    std::cout << "Sample data generated from normal distribution:\n"
            << plotVector(normalVector->as_flat_vector<double>()) << "\n";

    std::cout << "Sample data generated from user defined distribution:\n"
            << plotVector(userDefinedVector->as_flat_vector<double>()) << "\n";
    return 0;
}
